# 导入tushare
import tushare as ts

# 初始化pro接口
pro = ts.pro_api('31d3906685a1b0d31d52b4478cc1448a5ee4235abe9f55ab0a8943b9')

# 拉取数据
df = pro.daily(**{
    "ts_code": "600519.SH",
    "trade_date": "",
    "start_date": 20220414,
    "end_date": 20250414,
    "offset": "",
    "limit": ""
}, fields=[
    "ts_code",
    "trade_date",
    "open",
    "high",
    "low",
    "close",
    "pre_close",
    "change",
    "pct_chg",
    "vol",
    "amount"
])
print(df)

# 导入相关库
import numpy as np
import pandas as pd
import talib
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from sklearn.model_selection import GridSearchCV

# 2.简单衍生变量数据构造
df['close-open'] = (df['close'] - df['open']) / df['open']
df['high-low'] = (df['high'] - df['low']) / df['low']
df['pre_close'] = df['close'].shift(1)
df['price_change'] = df['close'] - df['pre_close']
df['p_change'] = (df['close'] - df['pre_close']) / df['pre_close'] * 100

# 3.移动平均线相关数据构造
df['MA5'] = df['close'].rolling(5).mean()
df['MA10'] = df['close'].rolling(10).mean()
df.dropna(inplace=True)

# 4.通过TA-Lib库构造衍生变量数据
df['RSI'] = talib.RSI(df['close'], timeperiod=12)
df['MOM'] = talib.MOM(df['close'], timeperiod=5)
df['EMA12'] = talib.EMA(df['close'], timeperiod=12) # 12日指数移动平均值
df['EMA26'] = talib.EMA(df['close'], timeperiod=26) # 26日指数移动平均值
df['MACD'], df['MACDsignal'], df['MACDhist'] = talib.MACD(df['close'], fastperiod=6, slowperiod=12, signalperiod=9)
df.dropna(inplace=True)

X = df[['close', 'close-open', 'MA5', 'MA10', 'high-low', 'RSI', 'MOM', 'EMA12', 'MACD', 'MACDsignal', 'MACDhist']]
y = np.where(df['price_change'].shift(-1) > 0, 1, -1)
X_length = X.shape[0]
split = int(X_length * 0.9)
X_train, X_test = X[:split], X[split:]
y_train, y_test = y[:split], y[split:]

from sklearn.model_selection import GridSearchCV
parameters = {'n_estimators': [5, 10, 20], 'max_depth': [2, 3, 4, 5, 6], 'min_samples_leaf': [5, 10, 20, 30]}
new_model = RandomForestClassifier(random_state=123)
grid_search = GridSearchCV(new_model, parameters, cv=6, scoring='accuracy')
grid_search.fit(X_train, y_train)
print("最佳参数:", grid_search.best_params_)

# 使用最佳模型进行预测
best_model = grid_search.best_estimator_
y_pred = best_model.predict(X_test)
y_pred_proba = best_model.predict_proba(X_test)

# 创建 DataFrame 显示预测值和实际值
a = pd.DataFrame()
a['预测值'] = list(y_pred)
a['实际值'] = list(y_test)
print(a.head())

# 创建 DataFrame 显示分类概率
a = pd.DataFrame(y_pred_proba, columns=['分类为-1的概率', '分类为1的概率'])
print(a.head())

# 计算准确率
accuracy = accuracy_score(y_pred, y_test)
print("准确率:", accuracy)

# 计算模型得分
score = best_model.score(X_test, y_test)
print("模型得分:", score)

# 获取特征重要性并排序
importances = best_model.feature_importances_
a = pd.DataFrame()
a['特征'] = X.columns
a['特征重要性'] = importances
a = a.sort_values('特征重要性', ascending=False)
print(a)

# 在测试数据上添加一列，预测收益
X_test['prediction'] = best_model.predict(X_test)

# 计算每天的股价变化率
X_test['p_change'] = (X_test['close'] - X_test['close'].shift(1)) / X_test['close'].shift(1)

# 计算累积收益率
X_test['origin'] = (X_test['p_change'] + 1).cumprod()

# 计算利用模型预测后的收益率
X_test['strategy'] = (X_test['prediction'].shift(1) * X_test['p_change'] + 1).cumprod()

X_test[['strategy', 'origin']].dropna().plot()
plt.gcf().autofmt_xdate()
plt.show()